from datetime import datetime
import akshare as ak
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np


def getData(code):
    data = ak.stock_zh_a_daily(symbol=code, adjust='qfq', start_date='20220104')
    return data.loc[:, ['date', 'close', 'high', 'low']]


def calcDonchianChannels(data, period=20):
    data['upperDon'] = data['high'].rolling(period).max()
    data['lowerDon'] = data['low'].rolling(period).min()
    data['midDon'] = (data['upperDon'] + data['lowerDon']) / 2
    return data


def midDonCrossOver(data, period=20, shorts=True):
    data['position'] = np.nan
    data['position'] = np.where(data['close'] > data['midDon'], 1, data['position'])

    if shorts:
        data['position'] = np.where(data['close'] < data['midDon'], -1, data['position'])
    else:
        data['position'] = np.where(data['close'] < data['midDon'], 0, data['position'])

    data['position'] = data['position'].ffill().fillna(0)

    return calcReturns(data)


def calcReturns(df):
    df['returns'] = df['close'] / df['close'].shift(1)
    df['log_returns'] = np.log(df['returns'])
    df['strat_returns'] = df['position'].shift(1) * df['returns']
    df['strat_log_returns'] = df['position'].shift(1) * df['log_returns']
    df['cum_returns'] = np.exp(df['log_returns'].cumsum()) - 1
    df['strat_cum_returns'] = np.exp(df['strat_log_returns'].cumsum()) - 1
    df['peak'] = df['cum_returns'].cummax()
    df['strat_peak'] = df['strat_cum_returns'].cummax()
    return df


def getStratStats(log_returns, risk_free_rate=0.02):
    stats = {}

    # Total Returns
    stats['tot_returns'] = np.exp(log_returns.sum()) - 1

    # Mean Annual Returns
    stats['annual_returns'] = np.exp(log_returns.mean() * 252) - 1

    # Annual Volatility
    stats['annual_volatility'] = log_returns.std() * np.sqrt(252)

    # Sortino Ratio
    annualized_downside = log_returns.loc[log_returns < 0].std() * np.sqrt(252)

    # Sharpe Ratio
    stats['sharpe_ratio'] = (stats['annual_returns'] - risk_free_rate) / stats['annual_volatility']

    # Max Drawdown
    cum_returns = log_returns.cumsum() - 1
    peak = cum_returns.cummax()
    drawdown = peak - cum_returns
    max_idx = drawdown.argmax()
    stats['max_drawdown'] = 1 - np.exp(cum_returns[max_idx] / np.exp(peak[max_idx]))

    # Max Drawdown Duration
    strat_dd = drawdown[drawdown == 0]
    strat_dd_diff = strat_dd.index[1:] - strat_dd.index[:-1]
    strat_dd_days = strat_dd_diff.map(lambda x: x.days).values
    strat_dd_days = np.hstack([strat_dd_days, (drawdown.index[-1] - strat_dd.index[-1]).days])
    stats['max_drawdown_duration'] = strat_dd_days.max()
    return {k: np.round(v, 4) if type(v) == np.float_ else v for k, v in stats.items()}


if __name__ == '__main__':
    code = 'sz000002'
    period = 20
    ticker = "XOM"
    res = getData(code)
    data = calcDonchianChannels(res, period)
    # colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
    #
    # plt.figure(figsize=(12, 8))
    # plt.plot(data["close"], label="Close")
    # plt.plot(data["upperDon"], label="Upper", c=colors[1])
    # plt.plot(data["lowerDon"], label="Lower", c=colors[4])
    # plt.plot(data["midDon"], label="Mid", c=colors[2], linestyle=":")
    # plt.fill_between(data.index, data["upperDon"], data["lowerDon"], alpha=0.3,
    #                  color=colors[1])
    #
    # plt.xlabel("Date")
    # plt.ylabel("Price ($)")
    # plt.title(f"Donchian Channels for {ticker}")
    # plt.xticks(rotation=45)
    # plt.legend()
    # plt.show()

    midDon = midDonCrossOver(data.copy(), 20, shorts=False)

    plt.figure(figsize=(12, 4))
    plt.plot(midDon["strat_cum_returns"] * 100, label="Mid Don X-Over")
    plt.plot(midDon["cum_returns"] * 100, label="Buy and Hold")
    plt.title("Cumulative Returns for Mid Donchian Cross-Over Strategy")
    plt.xlabel("Date")
    plt.ylabel("Returns (%)")
    plt.xticks(rotation=45)
    plt.legend()

    plt.show()

    stats = pd.DataFrame(getStratStats(midDon["log_returns"]),
                         index=["Buy and Hold"])
    stats = pd.concat([stats,
                       pd.DataFrame(getStratStats(midDon["strat_log_returns"]),
                                    index=["MidDon X-Over"])])
    print(stats)
